import idautils
import idaapi
import idc
import os

class Exporter:
    @staticmethod
    def check_decompiler():
        if not idaapi.init_hexrays_plugin():
            print("Hex-Rays Decompiler 未安装")
            return False
        return True

    @staticmethod
    def get_code(ea):
        try:
            cfunc = idaapi.decompile(ea)
            if cfunc is None:
                return None
            return "".join(idaapi.tag_remove(line.line) + "\n" for line in cfunc.get_pseudocode())
        except Exception as e:
            print(f"生成 {idc.get_func_name(ea)} 伪代码失败: {e}")
            return None

    @classmethod
    def export(cls, out_dir):
        if not cls.check_decompiler():
            return
        os.makedirs(out_dir, exist_ok=True)
        print("开始导出伪代码...")
        count = 0
        for ea in idautils.Functions():
            name = idc.get_func_name(ea)
            code = cls.get_code(ea)
            if code:
                safe_name = "".join(c if c.isalnum() or c == "_" else "_" for c in name)
                out_file = os.path.join(out_dir, f"{safe_name}.c")
                try:
                    with open(out_file, "w", encoding="utf-8") as f:
                        f.write(code)
                    count += 1
                    print(f"导出 {name} 到 {out_file}")
                except Exception as e:
                    print(f"写入 {out_file} 失败: {e}")
        print(f"完成，导出 {count} 个函数到 {out_dir}")

class Plugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    wanted_name = "ExportPseudocode"
    wanted_hotkey = "Alt+F8"

    def init(self):
        return idaapi.PLUGIN_OK if Exporter.check_decompiler() else idaapi.PLUGIN_SKIP

    def run(self, arg):
        out_dir = idaapi.ask_str("C:\\libc", 0, "输出目录：")
        if out_dir:
            Exporter.export(out_dir)

    def term(self):
        pass

def PLUGIN_ENTRY():
    return Plugin()